package scala.examples.fs;

import Syntax.{debug=>_, _};


class Typer {
  private def debug[t,a>: AnyRef](x: t): a =  Debug.tcheck.log(x)
  //private def error[t,a](x: t): a = Debug.tcheck.error(x)  
  
  type Label = String // TODO: refactor Select 
  type Env = LocMap[Type];

  import scoped._
  import util._
  
  class ScopedEnv extends Scoped[Env] { 
    // extend Gamma with binding 'name : tpe' in scope of 'scope'
    // NOTE: the execution of 'scope' should replace 'name' by the first argument 
    // passed to the block (so that free variables with the correct name are changed 
    // to a reference that will yield the correct type when looked up in env during the execution of scope)
    def extend[a](name: String)(tpe:Type)(scope: Ref=>a) :a= {
      val x = read.nextLoc
      val e = read.extend(tpe)
      //debug("extend: ("+x+", "+name+") : "+tpe)
      val r= withValue(e)(scope(Ref(x, name)))
      //debug("removed "+name)
      r
    }
  }    
    
  private val _env :ScopedEnv = new ScopedEnv
  //import _env.{read=>env,write=>env_=}
  def env  :Env = _env.read
  def env_= (v :Env) = _env.write(v)
  env = new Env
  
  private val _locked  = new ScopedLocks[AnyRef] // TODO split in Member's/ClassSig's/Term's ? (prob: lastField)
  def lock[A](l :AnyRef)(err :String)(scope: =>A): A = _locked.lock(l)(err)(scope)
  //import _locked.{read=>locked,write=>locked_=}
  def locked  :List[AnyRef] = _locked.read
  def locked_= (v :List[AnyRef]) = _locked.write(v)
  locked = Nil
    
  def isDecl(x :Member) = x.isInstanceOf[Declaration]

  def lastField(path :Term) :AnyRef/*Member \/ Term*/ = path match {
    case Select(p, a) => lookup(Single(p), a)
    case n: Name => n
    case n: Ref => n
  }
                                         

  // the \\uplus operator ---------------------------------------------

  def append(m: Member, ns: List[Member]): List[Member] = 
    if (ns exists (n => m.label == n.label)) ns else m :: ns;
  def append(ms: List[Member], ns: List[Member]): List[Member] = 
    (ms :\ ns) (append);
  def append(ms: List[List[Member]]): List[Member] = 
    (ms :\ List[Member]()) (append);

  // free variables ---------------------------------------------------

  private class FreeTest(x: Loc) extends Map {
    var result = false;
    override def term(e: Term): Term = e match {
      case Ref(y, _) if x == y => result = true; e
      case _ => super.term(e)
    }
  }
  
  private def contains(tp: Type, x: Loc): boolean = {
    val ft = new FreeTest(x); ft.typ(tp); ft.result
  }
  
  private def contains(m: Member, x: Loc): boolean = {
    val ft = new FreeTest(x); ft.member(m); ft.result
  }
  
  // TODO: correct? seems ad-hoc
  def checkNotContains(tp: Type, r: Ref): Type = {
    //debug("cnc: "+tp+" not: "+r+" "+contains(tp, r.x))

    if (contains(tp, r.x)) 
      tp match { 
	case Single(p) if(!(pathTypeOf(p) equals tp)) /* is this condition necessary ? */ =>
          checkNotContains(pathTypeOf(p), r);
                           
        case Inst(e) if(e.isPath && !(expand(Inst(e)) equals tp)) => 
          checkNotContains(expand(Inst(e)), r); 
          
	case _ => 
	  error("name " + r.name + " appears in term's type " + show.typ(tp));
      }
    else tp;
  }
  

  // Type assignment
  def pathTypeOf(p :Term) :Type = {
    assert(p.isPath)
    
    p match {
      case _ if (builtinTypeOf(p).isInstanceOf[Some[Type]]) => builtinTypeOf(p) match {case Some(x) => x}
      case TypedTerm(e, t) => if(!isSubType(typeOf(e), t)) error((show term e)+" does not conform to user-specified type "+(show typ t))
                              else t // NEW
      case Name(name) =>
        error("undefined: " + name); // typed names are substituted for references

      case Ref(loc, name) => // rule (Path-Var)
          env(loc) match {
           case Some(x) => x
           case None => error("undefined: " + name)
          }
          
          
      case Select(e1, l) => // rule (Path-Select)
        lookup(Single(e1), l) match {
          case ValSig(l, t, _) => t
          case DefSig(l, args, t, _) => FunType((args map { case Arg(_, t) => t } :List[Type]) ::: List(t))
        }
    }
  }
  
  def typeOf(e: Term): Type = { 
    debug("\n<typing term=\""+(show term e) +"\" case=\""+ showTree(e) + "\">\n");  // 
    
    val r= e match {
      case _ if (builtinTypeOf(e).isInstanceOf[Some[Type]]) => builtinTypeOf(e)  match {case Some(x) => x}
      
      case _ if (e.isPath) =>  // rule (Path)
          pathTypeOf(e)
          Single(e)            

      case TypedTerm(e, t) => if(!isSubType(typeOf(e), t)) error((show term e)+" does not conform to user-specified type "+(show typ t))
        else t  // NEW
      
      case e @ Select(t, a) if !e.isPath => // rule (Select)
        lookup(typeOf(t), a) match {
          case ValSig(a, tt, _) => tt
          case DefSig(a, args, rt, _) => FunType((args map { case Arg(_, at) => at } :List[Type]) ::: List(rt)) // NEW (from Method)
        }
 
      case Apply(fun, args) => // NEW (from Method)
        (if (fun.isPath) pathTypeOf(fun) else typeOf(fun)) match { 
          case FunType(ts) =>
            val as = args map (x => typeOf(x))

            assert(args.length <= ts.length-1) // ts contains return type
            
            if(List.forall2(as, ts)( 
                   (actual,formal) => isSubType(actual, formal))) {

              if(ts.length-args.length==1) 
                ts.last
              else
                FunType(ts.drop(args.length))
            }
            else {
              error("Argument mismatch in apply: "+(show term e)+" arguments: "+List.map2(as, ts)( 
                  (actual,formal) =>  ((show typ actual) + (if(isSubType(actual, formal)) " <: "  else " !<: ")+ (show typ formal))).mkString("", ", ", ""))
            }
        }
    
      case lv @ LocalVal(xname, tpe, value, scope) => // NEW (refactored from rule (New))
        val t = typeOf(value)
        
        val declt = tpe match {
          case Some(usert) =>  if (!isSubType(t, usert))
                                 error("illegal local value definition: "+show.term(lv))
                               else
                                 usert
          case None => t
        }
        
        wellformed(declt) 
        
        _env.extend(xname)(declt) (ref => { // env(x)=t (--> xname : t)
          val S = typeOf(subst(xname, ref) term scope) // Gamma[x :T] |- scope : S 

          checkNotContains(S, ref) // scope of local variable must be respected
        })

      case New(p) => // Refactored from rule (New)
        val t = Inst(p)
      
        wellformed(t)
      
        // TODO: what's the exact meaning of M_c
        val Pair(ms, self) = unfold(t);
        if (!ms.isEmpty && isDecl(ms.head)) {
	  error("cannot instantiate class with abstract member:\n  " + show.member(ms.head));
        }
      
        t
      
      // TODO: this is really just for convenience, not happy with it... find a more elegant solution
      case ClassDef(xname, c, scope) => // top-level class-definition, global scope
        // xname references a class, the type of a class is a class signature, which has a self type, namely, this class
        // replace the name of the class by the reference ref and extend Gamma with the corresponding type Class(c)
        val l = env.nextLoc
        val ref = Ref(l, xname)
        val declt = subst(xname, ref) typ Class(c)
        val newEnv = env.extend(declt)
        
        env = newEnv
        wellformed(declt)
        typeOf(subst(xname, ref) term scope)   
          
        //checkNotContains(S, ref) // scope of local variable must be respected, but these are global variables!
        
      case If(cond, then, els) =>
        if(!isSubType(typeOf(cond), Inst(Select(Name("$fs$"), "Boolean"))))
          error("if with non-boolean condition "+typeOf(cond))
          
        els match {
          case Some(els) => lub(then, els)
          case None => typeOf(then)
        }
    }
    
       
    debug("\n<res>"+(show typ r)+"\n</res>\n</typing>\n");
    r
  }
  
  // TODO: correct? "complete enough"?
  def lub(t1: Term, t2: Term): Type = lubMaybe(typeOf(t1),typeOf(t2)) match {
      case Some(t) => t 
      case _ =>  error("failed to compute lub of "+(show term t1)+": "+(show typ  typeOf(t1))+" and "+(show term t2) +": "+(show typ typeOf(t2)))
  }

  protected def lubMaybe(type1: Type, type2: Type): Option[Type] =
  {
    implicit def some[t](x: t) :Some[t] = Some(x)

    
    //debug("find lub of: "+(show typ type1)+" and "+(show typ type2))
    
    if(isSubType(type1,type2))
      type2
    else if(isSubType(type2,type1))
      type1
    else Pair(type1, type2) match {
      case Pair(Class(ClassSig(ps1, _,_, _)), Class(ClassSig(ps2, _,_, _))) => 
        (for(val p1 <- ps1; val p2 <- ps2) 
          yield lubMaybe(pathTypeOf(p1), pathTypeOf(p2))).find({case Some(t) => true case None => false}) match { case Some(ot) => ot case None => None }
          
      case _ => None
    }
  }


  
  def builtinTypeOf(e: Term): Option[Type] =  {
    implicit def some[t](x: t) :Some[t] = Some(x)
    e match {
        case Null => Single(Null) 
        
        case Name("$fs$") => BuiltinTypeFs
        case Name("unit")  => Inst(Select(Name("$fs$"), "Unit"))      
        case Name("bot")  => Inst(Select(Name("$fs$"), "Bottom"))      
        case True  => Inst(Select(Name("$fs$"), "Boolean"))
        case False  => Inst(Select(Name("$fs$"), "Boolean"))
        case StringLit(_) => Inst(Select(Name("$fs$"), "String"))
        //case class IntLit(i :Int) => Inst(Select(Name("$fs$"), "Int"))
        //case  Closure(args :List[Arg], t: Type, body: Term)  =>
        case _ => None
    }
    
  }

  private val stringSig = ClassSig(Nil, "this", Inst(Select(Name("$fs$"), "String")), Nil)
  private val booleanSig = ClassSig(Nil, "this", Inst(Select(Name("$fs$"), "Boolean")), Nil)
  private val unitSig = ClassSig(Nil, "this", Inst(Select(Name("$fs$"), "Unit")), Nil)
  private val bottomSig = ClassSig(Nil, "this", Inst(Select(Name("$fs$"), "Bottom")), Nil)
  
  // Expansion
  def unfold_debug (tp: Type): Pair[List[Member], String /*self*/] =  {
    debug("\n<unfold type=\""+(show typ tp) +"\"case=\""+ showTree(tp) + "\">\n") //self=\""+(show term self)+"\"
    val r = unfold(tp)

    for (val m <- r._1)  debug("<mem>"+(show member m)+"\n</mem>")
    debug("\n</unfold>\n")
    r
  }
  
  // compute the list of members in the type tp (name of this = self) in env
  def unfold(tp: Type): Pair[List[Member], String /*self*/] = tp match {
    case BuiltinTypeFs() => Pair(
      List( // members of the $fs$ object:
        ValSig("String", Class(stringSig), Some(ClassDef("this", stringSig, Name("unit")))),
        ValSig("Boolean", Class(booleanSig), Some(ClassDef("this", booleanSig, Name("unit")))),
        ValSig("Unit", Class(unitSig), Some(ClassDef("this", unitSig, Name("unit")))),
        ValSig("Bottom", Class(bottomSig), Some(ClassDef("this", bottomSig, Name("unit")))),
        DefDef("println", List(Arg("s", Inst(Select(Name("$fs$"), "String")))), Inst(Select(Name("$fs$"), "Unit")), Name("unit")),
        DefDef("print", List(Arg("s", Inst(Select(Name("$fs$"), "String")))), Inst(Select(Name("$fs$"), "Unit")), Name("unit")),
        DefDef("error", List(Arg("s", Inst(Select(Name("$fs$"), "String")))), Inst(Select(Name("$fs$"), "Bottom")), Name("bot"))
      ),"this$fs$")       

    case Class(ClassSig(ps, self, _, ms)) => // rule (Unfolds-SIGNATURE)
      val pss = ps map (p => unfold(Inst(p)) match {
        case Pair(members, self2) => subst(self2, Name(self)) members members 
        // NEW: rewrite members so that they all use our self name
      }) 
      
      // TODO: generic traverse
      val decls = pss map (.takeWhile(isDecl))
      val defs = pss map (.dropWhile(isDecl))
      Pair(append( decls ::: defs ::: List(ms.takeWhile(isDecl), ms.dropWhile(isDecl))), self)
             
    case Inst(Select(p, a)) =>
       lookup(Single(p), a) match { 
         case n @ TypeDef(a, t) => // rule (Unfolds-TYPE)
           lock(n)("cyclic inheritance involving type\n" + n){
             unfold(t)
           }
          
          case ValSig(a, _, Some(ClassDef(z, c @ ClassSig(ps, slf, _, ms), _))) => // rule (Unfolds-CLASS)
            lock(c)("cyclic inheritance involving class\n" + c){
              unfold(Class(c))
            }     
       }
    case Inst(p) =>  // NEW
      pathTypeOf(p) match {
        case cl @ Class(ClassSig(_, _,_, _)) => unfold(cl)
        case t @ Inst(Ref(l, _)) => error("unfold: "+env(l)+" t: "+t)
        case t => error("malformed type: " + show.typ(tp)+" (path type: "+(show typ t)+" )")
      }
  }
  

  
  // Membership (\ni)
  // lookup member l of type tp in environment env 
  def lookup(tp: Type, l: Label): Member = {
    def doLookup(tp: Type, newSelf: Term) = {
      val Pair(ms,self) = unfold(tp)
      ms find (m => m.label == l) match {
        case Some(m) =>  subst(self, newSelf) member m
        case _ => error("type " + show.typ(tp) + " does not contain member " + l);
      }
    }
    
    tp match {    
      case Single(p) => // rule (\ni-Singleton)
        val T = pathTypeOf(expand(p))
     
        lock(lastField(p))("circularity in membership") {
          doLookup(T, p)          
        }
     
      case _ => // rule (\ni-Other)
        _env.extend("this")(tp) (r => {
          val m = doLookup(tp, r) // just to check that doesn't contain self reference
              
          if (contains(m.signature, r.x))
            error("member " + (show member m) + " of type " + show.typ(tp) + 
                  "\n is not accessible because it contains a self reference `this'")
          else m
        })
    }
  }
  
  def lookup(p: Term): Member = p match {
    case Select(p, l) => lookup(typeOf(p), l)
  }
 
  // Well-Formedness  
  def wellformed_debug(t: Type): unit = { 
    debug("\n<wellformed case=\""+ showTree(t) + "\" type=\""+(show typ t) +"\">\n");  // 
    wellformed(t)
    debug("\n</wellformed>\n");
  }
  
  def wellformed(t: Type): unit = t match {
    case Single(p) => // WF-Singleton
       lock(lastField(p))("cycle involving type\n" + t){
         wellformed(pathTypeOf(p))
       }
    
    case Inst(Ref(loc, name)) => // NEW
      if (env(loc) == null) error("undefined: " + name)

    case Inst(Select(p, a)) => lookup(Single(p), a) match {
      case ValSig(a, Class(_), Some(_)) =>  // WF-Class
      case l @ TypeDef(a, t) =>  // WF-Type
         lock(l)("circularity in checking type alias"){
           wellformed(t)      
         }
      case TypeDcl(a, _) =>  // WF-Type'
    }
    
    case Class(c @ ClassSig(_, self, _, _)) => // WF-Signature
      _env.extend(self)(c.selfType) (selfRef => { // NEW: self: selfT instead of using ClassSig
        wellformed(c, selfRef) // NEW: can't easily substitute self in classSig here, as it's a bound variable (postponed to wellformed(classsig, ref)
      })
  }
    
   
  def wellformed_debug(m: Member, x: Term): unit = {
    debug("\n<wellformed-x member=\""+(show member m) +"\" case=\""+ showTree(m) + "\">\n");  // 
    wellformed(m, x)
  }
  
  def wellformed(m: Member, x: Term): unit = m match {

     case TypeDef(_, t) =>  wellformed(t) // WF-X-Type
     case TypeDcl(_, _) => // TODO: check upper bound
       
     case ValSig(l, t, Some(e @ ClassDef(xname, c @ ClassSig(_, self, _, _), _))) =>  // WF-X-Class
       _env.extend(xname)(Class(c)) (xref =>{
         _env.extend(self)(c.selfType) (selfRef => { // NEW: selfT instead of Inst(Select(x, l)) -- TODO
           wellformed(subst(xname, xref) classSig c, selfRef) 
           // TODO xname : Class(c) + subst xname --> needed in other places?
         })
       })
      
       if (!isSubType( typeOf(e), t)) // important! (see list.fs)
         error("illegal value definition: " + (show member m) + 
                "\n type of right-hand side " + (show typ typeOf(e)) + 
                "\n does not conform to declared type " + (show typ (t)))
                
     case ValSig(l, t, None) => // WF-X-Field
       wellformed(t) 
     
     case ValSig(l, t, Some(e)) => // WF-X-Field'
       wellformed(t) 
     
       if (!isSubType( typeOf(e), t))
         error("illegal value definition: " + (show member m) + 
                "\n type of right-hand side " + (show typ typeOf(e)) + 
                "\n does not conform to declared type " + (show typ (t)))
     
     case DefSig(l, args, t, o) => // WF-X-Method
       wellformed(t)
       for(val Arg(_, ati) <- args) {
         ati match { case Single(_) => error("argument types contained singleton type") case _ => } 

         wellformed(ati)
       }
       
       o match { 
         case Some(e) =>
           var env2 = env
           var body = e
          
           for(val Arg(xx, s) <- args) {
             val xRef = env2.nextLoc
             env2 = env2.extend(s)
             body = subst(xx, Ref(xRef, xx)) term body
           }
          
           _env.withValue(env2) {
             val bt = typeOf(body)
             if (!isSubType(bt, t)){
               error("illegal method definition: " + m + 
                     "\n body with type " + bt + 
                     "\n does not conform to declared type " + t)
             }
           }
         
         case None => 
       }
  }     

  
  def wellformed_debug(c: ClassSig, selfRef: Ref): unit = {
    debug("\n<wellformed-x type=\""+(show classSig c) +"\" case=\""+ showTree(c) + "\">\n");  // 
    wellformed(c, selfRef)
  }
  
  def wellformed(c: ClassSig, selfRef: Ref): unit = c match { // WF-X-Signature
    case ClassSig(ps, self, st, ms) => // NEW: selfT
      if(st!=null)
        wellformed(st)
      for(val mi <- ms) wellformed(subst(self, selfRef) member mi, selfRef)

      // check that parent types are well-formed, unfold them and substitute self-reference
      val nss = (for(val ti <- ps; {wellformed(Inst(ti));true}; 
                    val Pair(ni, selfi) = unfold(Inst(ti))) 
                   yield {subst(selfi, selfRef) members ni})

      // "The compatibility of a list of groups of members is defined in such a way that 
      //  a member declaration is always more precise than another declaration defined in 
      //  a previous group in this list"
      if(!foralli(nss){ (ni, i) =>
           nss.take(i).forall(prev => areSubTypes(ni ::: ms, prev))})
        error("invalid class signature" + c)
      
    case _ => error("invalid class signature" + c)
  }

  
  // Path Alias Expansion
  def expand(path :Term) :Term = pathTypeOf(path) match {
    case Single(q) => // rule (\simeq-Step)
       lock(lastField(path))("circularity in path alias expansion") {
          expand(q)      
       }
    case _ => path // rule (\simeq-Refl)
  }
  
  // Type Alias Expansion
  def expand(path :Type) :Type = path match {
    case Inst(Ref(loc,name)) => env(loc) match { case Some(x) => x case None => error("undefined: " + name) }// NEW
    case Inst(Select(p, a)) => lookup(Single(p), a) match {
        case n @ TypeDef(a, t) =>  // rule (\simeq-Type)
           lock(n)("circularity in type alias expansion") {
             expand(t)      
           }
        case TypeDcl(a, _) => path // \simeq-AbsType
        case ValSig(a, Class(_), Some(_)) => path // \simeq-Class
      }
    case FunType(ts) => FunType(ts map expand) // NEW
    case Class(_) => path // rule (\simeq-Signature)
    case Single(_) => path // rule (\simeq-Singleton)
  }
  

  // subtyping
  def isSubType_debug(t :Type, u :Type): Boolean = {val r = isSubType(t, u); debug("isSubType: "+(show typ t)+" and "+(show typ u)+" = "+r); r}

  def isSubType(t_ :Type, u_ :Type): Boolean = Pair(expand(t_), expand(u_)) match { // rule (<:-Unalias)
    case Pair(t, u) if t equals u => true // NEW <:-Refl
   
    case Pair(Single(p), Single(q)) => expand(p) equals expand(q) // rule (<:-Singleton-Right)

    
    case Pair(Single(p), u) => isSubType(pathTypeOf(expand(p)), u) // rule (<:-Singleton-left)

    case Pair(Inst(Select(Name("$fs$"), "Bottom")), _) => true // NEW
    case Pair(_ , Inst(Select(Name("$fs$"), "Any"))) =>  true // NEW
    
    case Pair(Inst(Select(p, a)), Inst(Select(q, b))) if (a equals b) => (expand(p) equals expand(q)) // rule (<:-Paths)

    case Pair(Inst(Select(p, a)), u @ Inst(Select(q, b))) if !(a equals b) =>   // rule (<:-Class)
      lookup(Single(p), a) match { 
        case ValSig(a, _, Some(ClassDef(_, c @ ClassSig(ps, _, _, _), _))) =>
          lock(c)("cyclic inheritance during subtyping involving class\n" + c){
            ps.exists(ti => isSubType(Inst(ti), u))
          }
      }

    //  rule (<:-Sig-Left):
    case Pair(Class(ClassSig(ps, _, _, _)), u @ Inst(Select(p, a))) => 
        ps.exists(ti => isSubType(Inst(ti), u))  
        
    case Pair(t, u @ Class(ClassSig(ps, self, _, ms))) if !t.isInstanceOf[Single] =>  // rule (<:-Sig-Right)
        ps.forall(ti => isSubType(t, Inst(ti))) && {
        val Pair(ns,nself) = unfold(t) 
        // TODO: is this subst ok?
        areSubTypes((subst(nself, Name(self)) members ns), ms)}   // NEW: rewrite members so that they all use the same self name
  
    case _ => false        
  }
  
//TODO: does this correctly capture the formalism (what is dom(M*)?)
  def areSubTypes(ms: List[Member], ns: List[Member]): Boolean =  
    ms.forall(m => if(ns exists (n => n.corresponds(m))) 
                      ns exists (n => isSubType(n, m))   
                   else   
                      true) 
  
  // member subtyping
  def isSubType(m: Member, n: Member): Boolean =  if(m equals n) true else Pair(m,n) match  {
    case Pair(TypeDef(a, t), TypeDef(b, t2)) if (a equals b) =>  isSubType(t,t2) // <:-member-type-2  -- NEW !!
    case Pair(ValSig(a, t, _), ValSig(b, t2, _)) if (a equals b) => isSubType(t,t2) // <:-member-field
    case Pair(DefSig(a, args, t, _), DefSig(b, args2, t2, _)) if (a equals b) => // <:-member-method
      isSubType(t, t2) &&
      List.forall2(args, args2)((a1,a2) => isSubType(a2.typ, a1.typ))
    case _ => false
  }
}

object util {
  def and(x: Boolean,y: Boolean) = x && y

  def foralli[a](list: List[a])(p: (a, Int) => Boolean): Boolean = {
    var these = list
    var i = 0
    while (!these.isEmpty) {
      if (!p(these.head, i)) return false
      these = these.tail
      i = i +1
    }
    true
  }
}

object scoped {
    class Scoped[T] { // aka DynamicVariable on wiki
      private val tl = new ThreadLocal
      def read :T = tl.get.asInstanceOf[T]
      def write(value :T) = tl.set(value)
      
      def withValue[A](value: T)(scope: =>A) = {
        val old = read
        try {
            write(value)
            scope
        } finally {
            write(old)
        }
      }
    }

 
    class ScopedLocks[t <:AnyRef] extends Scoped[List[t]] { 
      def lock[A](l :t)(err :String)(scope: =>A) :A= {
        if (read exists (l ==))
          error(err)
       
         withValue(read.::(l))(scope)
      }
    }
  }
